#!/bin/bash

# HDC 无线调试端口扫描器 - Shell版本
# 功能：扫描指定IP的所有端口，找到可以连接手机的HDC端口

# =========================== 配置参数 ===========================
SCAN_THREADS=500    # 并发扫描线程
HDC_THREADS=100     # HDC连接线程

# 端口扫描超时设置
TIMEOUT=0.02        # 端口检测超时时间

# 获取脚本所在目录的绝对路径
# 使用兼容的方法确保获取真实的绝对路径，避免符号链接问题
get_script_dir() {
    local source="${BASH_SOURCE[0]}"
    # 解析符号链接
    while [[ -L "$source" ]]; do
        local dir="$(cd -P "$(dirname "$source")" && pwd)"
        source="$(readlink "$source")"
        [[ $source != /* ]] && source="$dir/$source"
    done
    cd -P "$(dirname "$source")" && pwd
}

SCRIPT_DIR="$(get_script_dir)"
IP_HISTORY_FILE="$SCRIPT_DIR/ip_history.txt"

# =========================== 全局变量 ===========================
STOP_FLAG=0
FOUND_PORT=""
TEMP_DIR="/tmp/port_scanner_$$"

# 创建临时目录用于进程间通信
mkdir -p "$TEMP_DIR"

# =========================== 工具函数 ===========================

# 清理函数
cleanup() {
    echo ""
    echo "正在停止扫描..."
    STOP_FLAG=1
    
    # 终止所有子进程
    if [[ -f "$TEMP_DIR/pids" ]]; then
        while read -r pid; do
            kill "$pid" 2>/dev/null
        done < "$TEMP_DIR/pids"
    fi
    
    # 等待进程结束
    jobs -p | xargs -r kill 2>/dev/null
    wait 2>/dev/null
    
    # 清理临时文件
    rm -rf "$TEMP_DIR"
    
    echo "扫描结束，程序退出。"
    exit 0
}

# 设置信号处理
trap cleanup SIGINT SIGTERM EXIT

# 加载历史IP列表
load_ip_history() {
    if [[ -f "$IP_HISTORY_FILE" ]]; then
        # 读取所有非空IP地址，去重复
        grep -v '^$' "$IP_HISTORY_FILE" 2>/dev/null | awk '!seen[$0]++' || true
    fi
}

# 获取最近使用的IP
get_last_ip() {
    load_ip_history | head -1
}

# 保存IP到历史文件（最新的在最前面）
save_ip_to_history() {
    local new_ip="$1"
    local temp_file="$IP_HISTORY_FILE.tmp"
    
    # 创建临时文件，新IP放在第一行
    echo "$new_ip" > "$temp_file"
    
    # 添加其他不同的IP（最多保留9个历史IP）
    if [[ -f "$IP_HISTORY_FILE" ]]; then
        grep -v "^$new_ip$" "$IP_HISTORY_FILE" 2>/dev/null | head -9 >> "$temp_file" || true
    fi
    
    # 替换原文件
    mv "$temp_file" "$IP_HISTORY_FILE" 2>/dev/null || {
        echo "警告：无法保存IP到历史文件 $IP_HISTORY_FILE"
    }
}

# 显示IP选择菜单
show_ip_menu() {
    local ips=()
    local count=0
    
    # 读取历史IP
    while IFS= read -r ip; do
        if [[ -n "$ip" ]]; then
            ips+=("$ip")
            ((count++))
        fi
    done < <(load_ip_history)
    
    if [[ $count -eq 0 ]]; then
        return 1  # 没有历史IP
    fi
    
    # 显示前5条
    local display_count=$((count > 5 ? 5 : count))
    
    echo "" >&2
    echo "历史IP地址列表：" >&2
    echo "" >&2
    
    # 显示选项
    for ((i=0; i<display_count; i++)); do
        if [[ $i -eq 0 ]]; then
            echo "  $((i+1))) ${ips[i]} (最近使用)" >&2
        else
            echo "  $((i+1))) ${ips[i]}" >&2
        fi
    done
    
    echo "" >&2
    if [[ $count -gt 5 ]]; then
        echo "  f) 查看更多历史IP ($((count-5))个)" >&2
    fi
    echo "  n) 输入新的IP地址" >&2
    echo "" >&2
    
    while true; do
        echo -n "请选择 (1-$display_count$([ $count -gt 5 ] && echo ",f"),n) [回车选择1]: " >&2
        read -r choice
        
        # 直接回车选择第一个
        if [[ -z "$choice" ]]; then
            echo "${ips[0]}"
            return 0
        fi
        
        # 数字选择
        if [[ "$choice" =~ ^[0-9]+$ ]] && [[ $choice -ge 1 && $choice -le $display_count ]]; then
            echo "${ips[$((choice-1))]}"
            return 0
        fi
        
        # 查看更多
        if [[ "$choice" == "f" || "$choice" == "F" ]] && [[ $count -gt 5 ]]; then
            echo "" >&2
            echo "完整历史IP列表：" >&2
            echo "" >&2
            for ((i=0; i<count; i++)); do
                if [[ $i -eq 0 ]]; then
                    echo "  $((i+1))) ${ips[i]} (最近使用)" >&2
                else
                    echo "  $((i+1))) ${ips[i]}" >&2
                fi
            done
            echo "" >&2
            echo "  n) 输入新的IP地址" >&2
            echo "" >&2
            
            while true; do
                echo -n "请选择 (1-$count,n) [回车选择1]: " >&2
                read -r full_choice
                
                # 直接回车选择第一个
                if [[ -z "$full_choice" ]]; then
                    echo "${ips[0]}"
                    return 0
                fi
                
                # 数字选择
                if [[ "$full_choice" =~ ^[0-9]+$ ]] && [[ $full_choice -ge 1 && $full_choice -le $count ]]; then
                    echo "${ips[$((full_choice-1))]}"
                    return 0
                fi
                
                # 输入新IP
                if [[ "$full_choice" == "n" || "$full_choice" == "N" ]]; then
                    echo "" >&2
                    echo -n "请输入新的IP地址: " >&2
                    read -r new_ip
                    if [[ -n "$new_ip" ]]; then
                        echo "$new_ip"
                        return 0
                    fi
                fi
                
                echo "无效选项，请输入 1-$count 或 n" >&2
            done
        fi
        
        # 输入新IP
        if [[ "$choice" == "n" || "$choice" == "N" ]]; then
            echo "" >&2
            echo -n "请输入新的IP地址: " >&2
            read -r new_ip
            if [[ -n "$new_ip" ]]; then
                echo "$new_ip"
                return 0
            fi
        fi
        
        echo "无效选项，请输入 1-$display_count$([ $count -gt 5 ] && echo "，f")，n 或直接回车" >&2
    done
}

# 获取用户选择的IP
get_user_ip() {
    local last_ip
    last_ip=$(get_last_ip)
    
    if [[ -n "$last_ip" ]]; then
        echo "加载最近使用的IP: $last_ip" >&2
        echo "" >&2
        # 直接显示选择菜单
        local selected_ip
        selected_ip=$(show_ip_menu)
        echo "$selected_ip"
        return 0
    else
        # 没有历史IP，直接请求输入
        echo -n "请输入目标设备的IP地址: " >&2
        read -r ip
        echo "$ip"
        return 0
    fi
}

# 检查端口是否开放
check_port() {
    local ip="$1"
    local port="$2"
    
    # 检查停止标志
    [[ $STOP_FLAG -eq 1 ]] && return 1
    [[ -f "$TEMP_DIR/found_port" ]] && return 1
    
    # 使用nc命令检查端口（超时控制）
    if timeout "$TIMEOUT" nc -z "$ip" "$port" 2>/dev/null; then
        echo "$port" >> "$TEMP_DIR/open_ports"
        echo -e "\033[1;33m发现开放端口: $port\033[0m"
        
        # 异步执行hdc连接（限制并发数）
        while [[ $(jobs -r | grep "run_hdc_tconn" | wc -l) -ge $HDC_THREADS ]]; do
            sleep 0.01
            [[ $STOP_FLAG -eq 1 ]] && return 1
        done
        
        run_hdc_tconn "$ip" "$port" &
        echo $! >> "$TEMP_DIR/pids"
        return 0
    fi
    return 1
}

# 尝试HDC连接
run_hdc_tconn() {
    local ip="$1"
    local port="$2"
    
    # 检查停止标志和是否已找到端口
    [[ $STOP_FLAG -eq 1 ]] && return 1
    [[ -f "$TEMP_DIR/found_port" ]] && return 1
    
    # 执行hdc tconn命令
    local result
    result=$(timeout 3 hdc tconn "${ip}:${port}" 2>&1)
    local exit_code=$?
    
    if [[ $exit_code -eq 0 ]]; then
        if echo "$result" | grep -q "Connect OK"; then
            # 使用锁机制避免竞态条件
            (
                flock -x 200
                if [[ ! -f "$TEMP_DIR/found_port" ]]; then
                    echo -e "\033[1;32m$ip:$port HDC无线调试连接成功！\033[0m"
                    echo "$port" > "$TEMP_DIR/found_port"
                    STOP_FLAG=1
                fi
            ) 200>"$TEMP_DIR/lock"
            return 0
        elif echo "$result" | grep -q "Target is connected, repeat operation"; then
            (
                flock -x 200
                if [[ ! -f "$TEMP_DIR/found_port" ]]; then
                    echo -e "\033[1;32m$ip:$port HDC无线调试已连接（重复连接）\033[0m"
                    echo "$port" > "$TEMP_DIR/found_port"
                    STOP_FLAG=1
                fi
            ) 200>"$TEMP_DIR/lock"
            return 0
        fi
    fi
    
    return 1
}

# 扫描端口范围（带进度显示）
scan_ports_range() {
    local ip="$1"
    local start_port="$2"
    local end_port="$3"
    local range_name="$4"
    
    local total=$((end_port - start_port + 1))
    local count=0
    local progress_file="$TEMP_DIR/progress_${range_name// /_}"
    
    for ((port=start_port; port<=end_port; port++)); do
        # 检查停止条件
        [[ $STOP_FLAG -eq 1 ]] && break
        [[ -f "$TEMP_DIR/found_port" ]] && break
        
        # 控制并发数量
        while [[ $(jobs -r | wc -l) -ge $SCAN_THREADS ]]; do
            sleep 0.01
            [[ $STOP_FLAG -eq 1 ]] && break
        done
        
        check_port "$ip" "$port" &
        echo $! >> "$TEMP_DIR/pids"
        
        # 更新进度
        ((count++))
        local percent=$((count * 100 / total))
        echo "$percent" > "$progress_file"
    done
    
    # 标记完成
    echo "100" > "$progress_file"
}

# 显示扫描进度
show_scan_progress() {
    local ranges=("段1" "段2" "段3" "段4")
    
    while [[ $STOP_FLAG -eq 0 ]] && ! [[ -f "$TEMP_DIR/found_port" ]]; do
        local total_progress=0
        local completed_ranges=0
        
        for range in "${ranges[@]}"; do
            local progress_file="$TEMP_DIR/progress_${range// /_}"
            
            if [[ -f "$progress_file" ]]; then
                local percent
                percent=$(cat "$progress_file" 2>/dev/null || echo "0")
                if [[ "$percent" == "100" ]]; then
                    ((completed_ranges++))
                fi
                total_progress=$((total_progress + percent))
            fi
        done
        
        # 显示总体进度
        local overall_percent=$((total_progress / 4))
        echo -ne "\r\033[K\033[1;34m扫描进度: $overall_percent%\033[0m"
        
        # 全部完成则退出
        [[ $completed_ranges -eq 4 ]] && break
        
        sleep 0.3
    done
    
    echo ""  # 换行
}

# 开始扫描
start_scan() {
    local ip="$1"
    
    # 重置状态
    STOP_FLAG=0
    rm -f "$TEMP_DIR/found_port" "$TEMP_DIR/open_ports" "$TEMP_DIR/pids"
    touch "$TEMP_DIR/pids"
    
    echo ""
    echo "开始扫描 $ip 的端口..."
    echo ""
    
    # 记录开始时间
    local start_time
    start_time=$(date +%s)
    
    # 从25000开始扫描，分为4段多进程并发
    # 第一段：25000-32499
    scan_ports_range "$ip" 25000 32499 "段1" &
    local pid1=$!
    
    # 第二段：32500-39999
    scan_ports_range "$ip" 32500 39999 "段2" &
    local pid2=$!
    
    # 第三段：40000-52499
    scan_ports_range "$ip" 40000 52499 "段3" &
    local pid3=$!
    
    # 第四段：52500-65535
    scan_ports_range "$ip" 52500 65535 "段4" &
    local pid4=$!
    
    # 启动进度显示
    show_scan_progress &
    local progress_pid=$!
    
    # 等待扫描完成或找到端口
    while [[ $STOP_FLAG -eq 0 ]] && ! [[ -f "$TEMP_DIR/found_port" ]]; do
        # 检查扫描进程是否还在运行
        if ! kill -0 "$pid1" 2>/dev/null && ! kill -0 "$pid2" 2>/dev/null && \
           ! kill -0 "$pid3" 2>/dev/null && ! kill -0 "$pid4" 2>/dev/null; then
            break
        fi
        sleep 0.5
    done
    
    # 终止扫描进程
    kill "$pid1" "$pid2" "$pid3" "$pid4" "$progress_pid" 2>/dev/null
    wait "$pid1" "$pid2" "$pid3" "$pid4" "$progress_pid" 2>/dev/null
    
    # 计算耗时
    local end_time
    end_time=$(date +%s)
    local duration=$((end_time - start_time))
    
    show_scan_results "$ip" "$duration"
}

# 显示扫描结果
show_scan_results() {
    local ip="$1"
    local duration="$2"
    
    echo ""
    echo "=== 扫描结果 ==="
    
    # 显示结果
    if [[ -f "$TEMP_DIR/found_port" ]]; then
        local found_port
        found_port=$(cat "$TEMP_DIR/found_port")
        echo -e "\033[1;32m成功找到可连接的HDC端口: $found_port\033[0m"
        echo -e "   \033[1;36m连接命令: hdc tconn ${ip}:${found_port}\033[0m"
    else
        echo -e "\033[1;31m未找到可连接的HDC端口\033[0m"
        if [[ -f "$TEMP_DIR/open_ports" ]]; then
            local open_count
            open_count=$(wc -l < "$TEMP_DIR/open_ports" 2>/dev/null || echo 0)
            if [[ $open_count -gt 0 ]]; then
                echo "   发现 $open_count 个开放端口，但都无法建立HDC连接"
                echo "   开放的端口："
                head -10 "$TEMP_DIR/open_ports" | while read -r port; do
                    echo "     - $port"
                done
                [[ $open_count -gt 10 ]] && echo "     - ...还有 $((open_count - 10)) 个"
            fi
        else
            echo "   没有发现任何开放的端口"
        fi
    fi
    
    echo "   扫描耗时: ${duration}秒"
    echo ""
}

# 检查依赖命令
check_dependencies() {
    local missing_deps=()
    
    if ! command -v nc >/dev/null 2>&1; then
        missing_deps+=("nc (netcat)")
    fi
    
    if ! command -v hdc >/dev/null 2>&1; then
        missing_deps+=("hdc")
    fi
    
    if ! command -v timeout >/dev/null 2>&1; then
        missing_deps+=("timeout (coreutils)")
    fi
    
    if ! command -v flock >/dev/null 2>&1; then
        missing_deps+=("flock")
    fi
    
    if [[ ${#missing_deps[@]} -gt 0 ]]; then
        echo "❌ 错误：缺少以下依赖命令："
        printf '   - %s\n' "${missing_deps[@]}"
        echo ""
        echo "💡 在macOS上安装这些依赖："
        echo "   brew install netcat"
        echo "   brew install coreutils"
        echo "   # hdc需要从HarmonyOS开发工具包获取并添加到PATH"
        echo ""
        exit 1
    fi
}

# 验证IP地址格式
validate_ip() {
    local ip="$1"
    
    # 基本格式检查
    if ! [[ $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        return 1
    fi
    
    # 检查每个字段是否在0-255范围内
    IFS='.' read -ra ADDR <<< "$ip"
    for i in "${ADDR[@]}"; do
        if [[ $i -lt 0 || $i -gt 255 ]]; then
            return 1
        fi
    done
    
    return 0
}

# =========================== 主函数 ===========================
main() {
    echo "════════════════════════════════════════════════════════════"
    echo "HDC 无线调试端口扫描器 (Shell版本)"
    echo "   功能：自动扫描并连接HarmonyOS设备的无线调试端口"
    echo "════════════════════════════════════════════════════════════"
    
    # 检查依赖
    check_dependencies
    
    echo ""
    
    # 获取用户选择的IP
    local ip
    ip=$(get_user_ip)
    
    # 验证IP输入
    if [[ -z "$ip" ]]; then
        echo "❌ 未输入有效IP，程序退出。"
        exit 1
    fi
    
    # 验证IP格式
    if ! validate_ip "$ip"; then
        echo "❌ IP地址格式不正确: $ip"
        echo "   请输入正确的IPv4地址格式，如: 192.168.1.100"
        exit 1
    fi
    
    # 保存IP到历史
    save_ip_to_history "$ip"
    
    # 测试目标主机连通性（静默）
    ping -c 1 -W 1000 "$ip" >/dev/null 2>&1
    
    # 开始扫描
    start_scan "$ip"
}

# =========================== 脚本入口 ===========================
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi